import copy
import math

import common_modules

from DrawingAPIs.bar_drawer import *

def LayerWiseData(layerwise_file):
    layer_time = {}
    print(layerwise_file)
    with open(layerwise_file, newline='') as csvfile:
        rows = csv.reader(csvfile, delimiter=',')
        for row in rows:
            if row[2].lower() != "convolution":
                continue
            # print("Overall={0}, row[3]={1}".format(row, row[3]))
            layer_time[row[1]] = float(row[3][:-3])

    return layer_time

def OverallData(overall_file):
    net_time = {}
    with open(overall_file, newline='') as csvfile:
        rows = csv.reader(csvfile, delimiter=',')
        for row in rows:
            net_time[row[0]] = float(row[1][:-3])

    return net_time

def DrawLayerWiseSpeedup(devices, nets, parallel_data, serial_data, out_dir):
    # parallel_data, serial_data: {net model: {device name: {layer name: execution time}}}
    legend_labels = copy.deepcopy(devices)
    speedups = {}               # {net model: {device name: {layer name: speedup}}}
    max_speedup = 0.0
    max_speedup_dev = ''
    max_speedup_net = ''
    max_speedup_layer = ''
    min_speedup = 10.0
    min_speedup_dev = ''
    min_speedup_net = ''
    min_speedup_layer = ''
    for net in nets:
        devices = copy.deepcopy(legend_labels)
        speedups.setdefault(net, {})

        for dev in devices:
            if dev not in speedups[net].keys():
                speedups[net].setdefault(dev, {})

        # Compute the main reference speedup results.
        temp_dev = 'P100'
        for layer in parallel_data[net][temp_dev].keys():
            speedups[net][temp_dev][layer] = serial_data[net][temp_dev][layer] / parallel_data[net][temp_dev][layer]
            if speedups[net][temp_dev][layer] > max_speedup:
                max_speedup = speedups[net][temp_dev][layer]
                max_speedup_dev = temp_dev
                max_speedup_net = net
                max_speedup_layer = layer
            if speedups[net][temp_dev][layer] < min_speedup:
                min_speedup = speedups[net][temp_dev][layer]
                min_speedup_dev = temp_dev
                min_speedup_net = net
                min_speedup_layer = layer

        if len(speedups[net]['P100'].keys()) > 6:
            speedups[net][temp_dev] = dict([(k, speedups[net][temp_dev][k])
                                            for k in sorted(speedups[net][temp_dev],
                                                            key=speedups[net][temp_dev].get,
                                                            reverse=True)])

        devices.remove('P100')
        for dev in devices:
            for layer in speedups[net]['P100'].keys():
                speedups[net][dev][layer] = serial_data[net][dev][layer] / parallel_data[net][dev][layer]
                if speedups[net][dev][layer] > max_speedup:
                    max_speedup = speedups[net][dev][layer]
                    max_speedup_dev = dev
                    max_speedup_net = net
                    max_speedup_layer = layer
                if speedups[net][dev][layer] < min_speedup:
                    min_speedup = speedups[net][dev][layer]
                    min_speedup_dev = dev
                    min_speedup_net = net
                    min_speedup_layer = layer

    # Print the maximum speedup achieved on a single layer.
    print("Maximum speedup achieved on a single network layer: {0}--{1} on {2} = {3}".format(max_speedup_net,
                                                                                             max_speedup_layer,
                                                                                             max_speedup_dev,
                                                                                             max_speedup))
    print("Minimum speedup achieved on a single network layer: {0}--{1} on {2} = {3}".format(min_speedup_net,
                                                                                             min_speedup_layer,
                                                                                             min_speedup_dev,
                                                                                             min_speedup))
    net_layers_used = {}
    for net in nets:
        x_ticklabels = []
        speedup_values = []         # [[device1_values], [device2_values], ..., [devicen_values]]
        if len(speedups[net]['P100'].keys()) > 6:
            # If the number of convolution layers is larger than 6, we only pick the first 6 layers.
            x_ticklabels = list(speedups[net]['P100'].keys())[:6]
        else:
            x_ticklabels = list(speedups[net]['P100'].keys())

        net_layers_used[net] = copy.deepcopy(x_ticklabels)
        bar_width = 0.45
        x_ticks = np.arange(2, 2 * (len(x_ticklabels) + 1), 2)
        y_maxval = 0.0
        for dev in legend_labels:
            temp_speedup_val = []
            for layer in x_ticklabels:
                temp_speedup_val.append(speedups[net][dev][layer])
            y_maxval = max(y_maxval, max(temp_speedup_val))
            speedup_values.append(temp_speedup_val)
        # Draw Net-specific diagrams.
        if y_maxval > (math.ceil(y_maxval) - 0.5):
            y_maxval = math.ceil(y_maxval)
        else:
            y_maxval = math.ceil(y_maxval) - 0.5

        for i in range(0, len(x_ticklabels), 1):
            x_ticklabels[i] = 'conv' + str(i + 1)

        # print("Drawing the speedup figure of the {0} model...".format(net))
        DrawBarFig(x_ticks, x_ticklabels, speedup_values, legend_labels, '', 'Speedup',
                   bar_width, x_min=0.8, x_max=(x_ticks[-1] + 1), y_min=0, y_max=y_maxval,
                   folder=out_dir + '/Speedup/IntraLayer', filename='exp_speedup_' + net, allow_legend=True)
    return net_layers_used

def DrawOverallSpeedup(devices, nets, parallel_data, serial_data, out_dir):
    # parallel_data, serial_data: {device: {net: execution time}}
    speedups = []
    y_maxval = 0.0
    for dev in devices:
        dev_speedup = []
        for net in nets:
            dev_speedup.append(serial_data[dev][net] / parallel_data[dev][net])
        y_maxval = max(y_maxval, max(dev_speedup))
        speedups.append(dev_speedup)

    print("Maximum speedup is: {0}".format(y_maxval))
    bar_width = 0.45
    x_ticks = np.arange(2, 2 * (len(nets) + 1), 2)
    if y_maxval > (math.ceil(y_maxval) - 0.5):
        y_maxval = math.ceil(y_maxval)
    else:
        y_maxval = math.ceil(y_maxval) - 0.5

    print("Drawing NNForwardSpeedup figure...")
    DrawBarFig(x_ticks, x_ticklabels=nets, y_values=speedups, legend_labels=devices,
               x_label="Neural Network Model", y_label='Speedup', bar_width=bar_width, x_min=0.8,
               x_max=(x_ticks[-1] + 1), y_min=0, y_max=y_maxval,
               folder=out_dir + '/IntraLayer/Speedup', filename='NNForwardSpeedup', allow_legend=True)
    return

def DrawIntraLayerSpeedup(args, devices):
    # Parse serial execution time.
    nets = []
    serial_layerwise_time = {}
    serial_overall_time = {}
    for dev in devices:
        curr_root = args.root[0] + '/' + dev + '/' + args.serial[0]
        serial_overall_time.setdefault(dev, {})
        print("Serial folder: {0}".format(curr_root))
        for net in os.listdir(curr_root):
            if net.lower() == "googlenet":
                continue
            serial_layerwise_time.setdefault(net, {})
            serial_layerwise_time[net].setdefault(dev, {})

            if net not in nets:
                nets.append(net)
            layerwise_log_file = curr_root + '/' + net + '/' + net + 'LayerForwardTime.csv'
            overall_log_file = curr_root + '/' + net + '/' + net + 'OverAllForward.csv'

            serial_layerwise_time[net][dev].update(LayerWiseData(layerwise_log_file))
            serial_overall_time[dev].update(OverallData(overall_log_file))

    # Parse HGP4NN execution time.
    hgp4nn_layerwise_time = {}
    hgp4nn_overall_time = {}
    for dev in devices:
        curr_root = args.root[0] + '/' + dev + '/' + args.hgp4nn[0]
        print("HGP4NN folder: {0}".format(curr_root))
        hgp4nn_overall_time.setdefault(dev, {})
        for net in nets:
            hgp4nn_layerwise_time.setdefault(net, {})
            hgp4nn_layerwise_time[net].setdefault(dev, {})

            layerwise_log_file = curr_root + '/' + net + '/LayerWise/' + net + 'LayerForwardTime-PROF.csv'
            overall_log_file = curr_root + '/' + net + '/LayerWise/' + net + 'OverAllForward-PROF.csv'

            hgp4nn_layerwise_time[net][dev].update(LayerWiseData(layerwise_log_file))
            hgp4nn_overall_time[dev].update(OverallData(overall_log_file))

    print('\n')
    net_layer_used = DrawLayerWiseSpeedup(devices, nets, hgp4nn_layerwise_time, serial_layerwise_time, args.output[0])
    DrawOverallSpeedup(devices, nets, hgp4nn_overall_time, serial_overall_time, args.output[0])

    return net_layer_used
